#!/usr/bin/ruby

# Summary: monitor who is eating your bandwidth
# WWW:     http://code.google.com/p/trafmonitor/
# Author:  tmarciniak@gmail.com
# License: New BSD

require 'pcaplet'
require 'lib/trafmonitor'

include LibTrafmonitor

def pid_handler( pidfile_path )
    $LOG.debug do "Checking for pidfile existence." end if $DEBUG

    if  File.exist?( pidfile_path )
        $LOG.fatal do "Is trafmonitor already running? Check #{pidfile_path}. Terminating." end
        exit 1
    else
        $LOG.debug do "Writing pidfile #{pidfile_path}." end if $DEBUG
        File.open( pidfile_path, 'w+' ) { |fh| fh.puts Process.pid }
    end
end

def main
    pidfile   = $config.pidfile
    filterexp = $config.filter
    stats     = Stats.new
    sniffer   = Pcaplet.new( '-i eth0 -s 1500' )
    my_filter = Pcap::Filter.new( filterexp, sniffer.capture )
    
    pid_handler( pidfile )
    sniffer.add_filter( my_filter )

    $LOG.debug do 'Starting packet sniffer loop.' end if $DEBUG
    
    sniffer.each_packet { |p|
		stats.account( p.ip_dst.to_num_s, p.size.to_i )
    
        Signal.trap( "USR1" ) do
            $LOG.debug do 'Stats dumped with SIGUSR1' end if $DEBUG
            stats.to_log
        end
    
        Signal.trap( "USR2" ) do
            $LOG.debug do 'Stats cleared with SIGUSR2' end if $DEBUG
            stats.reset
        end

        Signal.trap( "INT" ) do
            $LOG.info do 'Caught interrupt. Dumping stats and terminating.' end
            stats.to_log
            begin
                File.delete( $config.pidfile )
            rescue SystemCallError
            end
            exit 0
        end
    }
end

$config = Config.new( '/root/.trafmonitorrc' )
$LOG    = Logger.new( $config.logfile )

if $DEBUG
    $LOG.debug do 'Debugging session started' end
    main()
else
    pid = fork do
        main()
    end
    
    $LOG.info do "trafmonitor daemon started with pid #{pid}." end
    Process.detach( pid )
end
